from typing import TypedDict, Dict, Any, Annotated

from langgraph.constants import END
from langgraph.graph import StateGraph

from owl_ai.domain.agent_config_entity import AgentConfigEntity
from owl_ai.graph.graph import BaseGraph


def update_node_params(state: Dict[str, Any], node_output: Any) -> Dict[str, Any]:
    """
    自定义 reducer 函数，用于更新 node_params 字段
    """
    # 如果 node_params 不存在，则初始化为空字典
    if "node_params" not in state:
        state["node_params"] = {}

    if node_output:
        # 将节点的输出添加到 node_params 中
        state["node_params"].update(node_output)
    return state["node_params"]


class WorkFlowState(TypedDict):
    node_params: Annotated[dict, update_node_params]


class WorkFlowGraph(BaseGraph):

    @classmethod
    def find_node_config(cls, node_config_list: list[dict], node_name: str):
        for node_config in node_config_list:
            if node_config.get("nodeName") == node_name:
                return node_config
        return None

    @classmethod
    def compile(cls, graph_config: AgentConfigEntity):
        """
        工作流图编译
        Args:
            graph_config: 图配置

        Returns: 编译后的langgraph图对象
        """
        graph = StateGraph(WorkFlowState)

        node_config_list = graph_config.config
        for node_config in node_config_list:
            node_type = node_config.get("nodeType")

            # END节点不需要处理
            if END == node_type:
                continue
            node_name = node_config.get("nodeName")
            if node_type == "llm":
                from owl_ai.graph.flow.chat_llm_graph_node import ChatLLMGraphNode
                graph.add_node(node_name, ChatLLMGraphNode(node_config))
            if node_type == "rag":
                from owl_ai.graph.flow.rag_retrieval_graph_node import RagRetrievalGraphNode
                graph.add_node(node_name, RagRetrievalGraphNode(node_config))
            if node_type == "agent":
                from owl_ai.graph.flow.graph_agent_node import AgentGraphNode
                graph.add_node(node_name, AgentGraphNode(node_config))

            next_nodes = node_config.get("nextNodes")
            for next_node in next_nodes:
                next_node_config = cls.find_node_config(node_config_list, next_node)
                next_node_type = next_node_config.get("nodeType")
                if next_node_type == "condition":
                    from owl_ai.graph.flow.graph_router import GraphConditionRouter
                    graph.add_node(next_node, GraphConditionRouter(next_node_config))
                    graph.add_edge(node_name, next_node)
                else:
                    graph.add_edge(node_name, next_node)

        return graph.compile()
